import sys, os
sys.path.append(f"{os.path.dirname(__file__)}")
sys.path.append(f"{os.path.dirname(__file__)}/../repESP_old")

import resp_helpers, cube_helpers, graphs

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import numpy as np

import argparse


def main():

    help_description = """
        Plot fitting or cube points in 2 or 3D colored by field values

        In order to avoid having a cluttered plot the user is recommended to
        specify a 'slice' of the points to be plotted. Only points within a given
        distance from the slicing plane will be shown. To enable this option
        see section 'optional arguments related to slicing' section below.
        """

    parser = argparse.ArgumentParser(
        description=help_description,
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument("points_file",
                        help="""the file containing the points around a molecule as
                        a cube file or fitting points in one of the two .esp
                        formats:

                        1. generated by running a Gaussian Pop calculation (for
                        a selected ESP-based method) with IOp(6/50)=1.

                        2. in the Antechamber format generated by `repESP` with
                        either the `rep_esp.py` or `field_diff.py` script."""
                        # TODO: Must be able to read in the identities from some
                        # other file or from a list
                        """
                        LIMITATION:
                        This format loses information about the atomic labels,
                        which will be plotted as 'XX'. As a workaround to decipher
                        what atoms the labels correspond to, re-run this script
                        with identical options but on a Gaussian-generated .esp
                        file.""",
                        metavar="POINTS_FILE")

    parser.add_argument("dimension",
                        help="""The number of dimensions for the plot. Note that 2D
                        plots require specifying a slicing plane with the
                        '--slice_dist' option.""",
                        choices=[2, 3], type=int, metavar="DIMENSION")

    parser.add_argument("--limits",
                        help="""limits on plot axes in angstroms. The same limits
                        will be applied to all axes, non-square/cubic plots are
                        currently not supported.""",
                        nargs=2, type=float,
                        default=(-5, 5), metavar=("LOWER", "UPPER"))

    mutex_colors = parser.add_mutually_exclusive_group()

    mutex_colors.add_argument("--color_range",
                              help="Specify the lower and upper limits for the "
                              "color range for field values at fitting points. If "
                              "this option is not specified, the limits will be "
                              "calculated automatically based on all data points, "
                              "not only the plotted slice of points.",
                              nargs=2, type=float, metavar=("LOWER", "UPPER"))

    mutex_colors.add_argument("--sym_color_range",
                              help="Use a symmetric automatic color range for the "
                              "colors of field values at fitting points.",
                              action="store_true")

    parser.add_argument("-o", "--output",
                        help="""output file name. The extension will determine
                        output format (.pdf recommended). If you don't want the
                        plot to be saved but only displayed in an interactive
                        window, leave this option unset (recommended for 3D
                        graphs).""")

    slicing_group = parser.add_argument_group(
        title="optional arguments related to slicing",
        description="""Switch slicing on by passing a value to '--slice_dist' and
        specifying the slicing plain with either '--slice_atoms' or '--slice_eqn'.
        Atoms lying on the slicing plane will be shown in a red box, atoms within
        0.5 A from the plane in a dashed red box and atoms outside that distance
        will be shown in grey.""")

    slicing_group.add_argument(
        "--slice_dist",
        help="The distance in angstrom from the slicing plane within which points "
        "are to be plotted. If all points are to be plotted, specify a very high "
        "number e.g. 1000",
        type=float, metavar=("DIST"))

    mutex_slicing_group = slicing_group.add_mutually_exclusive_group()

    mutex_slicing_group.add_argument(
        "--slice_atoms",
        help="The atoms through which the slicing plane should pass",
        nargs=3, type=int, metavar=("ATOM1_LABEL", "ATOM2_LABEL", "ATOM3_LABEL"))

    mutex_slicing_group.add_argument(
        "--slice_eqn",
        help="""The equation for the slicing plane specified with a list of
        parameters of the following plane equation:
        Ax + By + Cz + D = 0
        """,
        nargs=4, type=float, metavar=("A", "B", "C", "D"))

    args = parser.parse_args()

    if args.slice_dist is None and (args.slice_atoms or args.slice_eqn):
        if args.slice_atoms:
            wrong_option = 'slice_atoms'
        elif args.slice_eqn:
            wrong_option = 'slice_eqn'
        raise ValueError("The option '--{0}' should only be specified when the "
                         "'--slice_dist' option is set.".format(wrong_option))

    if args.slice_dist is not None and not (args.slice_atoms or args.slice_eqn):
        raise ValueError("The option '--slice_dist' requires either '--slice_eqn' "
                         "or '--slice_atoms' to be set.")

    if args.output is not None and os.path.exists(args.output):
        raise FileExistsError("Output file exists: " + args.output)

    if args.dimension == 2 and not args.slice_dist:
        raise ValueError("A 2D plot requires a slicing plane to be requested with "
                         "the '--slice_dist' option.")

    filetype = args.points_file[-4:]
    if filetype not in ['.cub', '.esp']:
        print("The fields must have the extension .cub or .esp")
        sys.exit(1)

    if filetype == '.cub':
        g = cube_helpers.Cube(args.points_file)
    elif filetype == '.esp':
        g = resp_helpers.G09_esp(args.points_file)

    # TODO: The colour span is calculated based on all values, but if only a slice
    # is plotted, the shown points may occupy only a narrow range.
    if args.color_range:
        color_span = args.color_range
    elif args.sym_color_range:
        color_limit = max(abs(np.nanmin(g.field.get_values())),
                          abs(np.nanmax(g.field.get_values())))
        color_span = [-color_limit, color_limit]
    else:
        color_span = [min(g.field.get_values()), max(g.field.get_values())]

    if args.slice_atoms:
        plane_eqn = graphs.plane_through_atoms(g.molecule, *args.slice_atoms)
    else:
        plane_eqn = args.slice_eqn

    graphs.plot_points(
        g.field, args.dimension, molecule=g.molecule,
        plane_eqn=plane_eqn,
        dist_thresh=args.slice_dist,
        axes_limits=[args.limits]*args.dimension,
        color_span=color_span,
        save_to=args.output)
